/*
 * Copyright (C) 2010 Sinapsi Spa
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.sinapsi.memtracker.maven.tracker;

import static java.io.File.separator;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryPoolMXBean;
import java.lang.management.MemoryType;
import java.lang.management.MemoryUsage;
import java.net.MalformedURLException;
import java.util.LinkedList;
import java.util.List;

import javax.management.AttributeNotFoundException;
import javax.management.InstanceNotFoundException;
import javax.management.MBeanException;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import javax.management.ReflectionException;
import javax.management.openmbean.CompositeData;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.thoughtworks.xstream.XStream;
import com.thoughtworks.xstream.io.xml.DomDriver;

public class MemoryTracker {

    public static final String DEFAULT_DIR = "target/mem-tracker/";
    private static final String DEFAULT_SERVER_PORT = "10088";
    private static final String DEFAULT_SERVER_HOST = "localhost";
    private static Logger log = LoggerFactory.getLogger(MemoryTracker.class);

    private String url;

    private String outputDir = DEFAULT_DIR;

    private String outputFile = "memory-usage.xml";

    public MemoryTracker() {
        this.url = "service:jmx:rmi:///jndi/rmi://" + DEFAULT_SERVER_HOST + ":" + DEFAULT_SERVER_PORT + "/jmxrmi";
    }

    public MemoryTracker(String url) {
        this.url = url;
    }

    public MemoryTracker(String host, String port) {
        this.url = "service:jmx:rmi:///jndi/rmi://" + host + ":" + port + "/jmxrmi";
    }

    public void saveRemote() throws Exception {
        List<Memory> memories = getRemoteMemories();
        save(memories);
        log.info("tracked memory usage for server [{}] on port [{}]:", DEFAULT_SERVER_HOST, DEFAULT_SERVER_PORT);
        for (Memory memory : memories) {
            log.info("{}", memory);
        }
    }

    protected List<Memory> getRemoteMemories() throws Exception {
        List<Memory> memories = new LinkedList<Memory>();
        JMXConnector connection = null;
        try {
            connection = getConnection();
            List<MemoryPoolMXBean> pools = ManagementFactory.getMemoryPoolMXBeans();
            for (MemoryPoolMXBean pool : pools) {
                memories.add(new Memory(pool.getName(), pool.getType(), getRemoteUsage(connection, pool)));
            }
        } finally {
            if (connection != null) {
                connection.close();
            }
        }
        return memories;
    }

    protected CompositeData getRemoteUsage(JMXConnector connection, MemoryPoolMXBean pool) throws MalformedObjectNameException, MBeanException,
            AttributeNotFoundException, InstanceNotFoundException, ReflectionException, IOException {
        ObjectName name = new ObjectName("java.lang:type=MemoryPool,name=" + pool.getName());
        CompositeData peak = (CompositeData) connection.getMBeanServerConnection().getAttribute(name, "PeakUsage");
        return peak;
    }

    protected void save(List<Memory> memories) throws FileNotFoundException, IOException {
        new File(outputDir).mkdir();
        XStream xstream = buildXStream();
        FileOutputStream out = new FileOutputStream(outputDir + separator + outputFile);
        xstream.toXML(memories, out);
        out.flush();
        out.close();
    }

    protected JMXConnector getConnection() throws IOException, MalformedURLException {
        return JMXConnectorFactory.connect(new JMXServiceURL(url), null);
    }

    @SuppressWarnings("unchecked")
    public static List<Memory> load(String file) {
        return (List<Memory>) buildXStream().fromXML(file);
    }

    private static XStream buildXStream() {
        XStream xstream = new XStream(new DomDriver());
        xstream.alias("memory", Memory.class);
        xstream.alias("usage", MemoryUsage.class);
        xstream.alias("type", MemoryType.class);
        return xstream;
    }

    public void setOutput(String output) {
        this.outputDir = new File(output).getParent();
        this.outputFile = new File(output).getName();
    }
}
